# scidb.py 
import re
import MySQLdb, pymym
from numpy import *
import numpy as np
from pylab import find
from itertools import izip

UNKNOWN = 0
PULLED = 1
MODIFIED = 2

#################################################################################
# MySQL commands
MYSQL_ADD_COL = 'alter table %s add column %s %s null' 
MYSQL_UPDATE_ROW = "update %s SET %s='%s' WHERE %s=%d"
MYSQL_DROP_COL = 'alter table %s drop column %s'

#################################################################################
# Helper Functions
def mysql_get_entry(c, col_id, table_name, row_id, id='id'):
    if not iterable(row_id):
        row_id = [row_id]
    id_set = '('+ ', '.join([str(x) for x in row_id]) +')'
    c.execute("select %s from %s where %s IN %s" % (col_id, table_name, id, id_set))
    
    # mYm datatype conversion
    vals = c.fetchall()
    conv_vals = [ pymym.unpack(val[0]) for val in vals ]
    #conv_vals = []
    # for val in vals:
    #     val = val[0]
    #     pymym.unpack(val[0])
    #     # if pymym.is_mym_data(val):
    #     #     conv_vals.append( pymym.unpack(val) )
    #     # else:
    #     #     conv_vals.append(val)
    return conv_vals 

def execute_silent(c, cmd, args):
    try:
        c.execute(cmd % args)
    except:
        pass

def check_tbl_for_value(val):
    """ Return function to check table for value 
    """
    return lambda s, table: table[s, val] == '' or table[s,val] == None
#################################################################################
# Lazy Table 
class lazy_table():
    # TODO setup a config file for the default values..
    def __init__(self, table_name, db, config, id=None, cache_data=True, 
            verbose=False, name_unique=True):
        """
        config: dic t
            contain s fields 'USER', 'PASSWORD', 'HOST', 'PORT'
        """
        host = config['HOST']
        user = config['USER']
        passwd = config['PASSWORD']

        self.db_name = db
        self.table_name = table_name
        self.cache_data = cache_data
        
        ## establish database connection, or create DB if it does not exist
        try:
            self.conn = MySQLdb.connect(host=host, user=user, passwd=passwd, 
                db=db)
        except MySQLdb.OperationalError as (mysql_err_code, err):
            self.conn = MySQLdb.connect(host=host, user=user, passwd=passwd)

            if mysql_err_code == 1049: # database doesn't exist
                c = self.conn.cursor()
                c.execute('create database %s' % db)
                c.execute('use %s' % db)

        # see if tbl_name already exists in database
        c = self.conn.cursor()
        c.execute('show tables')
        table_names = [x[0] for x in c.fetchall()]

        if not table_name in table_names: # new table
            name_unique_str = 'unique' if name_unique else 'null'
            if verbose:
                print "Creating new table in database.."
            create_table_sql_cmd = """create table if not exists %s 
            (id INT NOT NULL AUTO_INCREMENT, name varchar(200) %s, 
            primary key (id))""" % (self.table_name, name_unique_str) 
            c.execute( create_table_sql_cmd )
        self.table_name = table_name

        ## pull column names and number of rows from table 
        self.update_column_names()

        if id == None:
            if not 'id' in self.col_names:
                django_inheritance_ids = filter( lambda x: re.match('.*?_ptr_id', x), self.col_names ) 
                if len(django_inheritance_ids) == 1:
                    id = django_inheritance_ids[0]
            else:
                id = 'id'
        self.id = id
        
        # check that table has columns for id, name
        if id in self.col_names and 'name' in self.col_names:
            # check that id and name columns are of the proper data type
            id_descr = filter(lambda x: x[0]==id, self.table_descr)[0]
            name_descr = filter(lambda x: x[0]=='name', self.table_descr)[0]

            id_valid_type = 'int' in id_descr[1] # and 'auto_increment' in id_descr
            name_valid_type = 'varchar' in name_descr[1]
            if not id_valid_type: 
                raise Exception("Invalid column type for 'id': %s" % id_descr[1])
            if not name_valid_type:
                raise Exception("Invalid column type for 'name'")

            self.name_unique = filter(lambda x: x[0]=='name', self.table_descr)[0][3] == 'UNI'
        elif id in self.col_names:
            id_descr = filter(lambda x: x[0]==id, self.table_descr)[0]
            id_valid_type = 'int' in id_descr[1] # and 'auto_increment' in id_descr
            if not id_valid_type: 
                raise Exception("Invalid column type for 'id': %s" % id_descr[1])

            self.name_unique = False 
        else:
            raise Exception("Table does not contain 'id' or 'name' columns")

        # cache rows
        n_rows = int( c.execute('select %s from %s' % (id, table_name)) ) 
        self.n_rows = n_rows
        self.table_vals = {}
        for col in self.col_names:
            self.table_vals[col] = [None]*n_rows

        self.status_table = None
        self.update_keys()

    def update_column_names(self):
        """ Pull column names and number of rows from table
        """
        c = self.conn.cursor()
        c.execute('describe %s' % self.table_name)
        self.table_descr = c.fetchall()
        self.col_names = [x[0] for x in self.table_descr]

    def update_keys(self):
        c = self.conn.cursor()
        if self.name_unique:
            c.execute('select %s, name from %s' % (self.id, self.table_name))

            name_id_pairs = c.fetchall() 
            self.names = [x[1] for x in name_id_pairs ]
            self.ids = [x[0] for x in name_id_pairs]

            # insert names and ids into table values  
            self.table_vals[self.id] = self.ids
            self.table_vals['name'] = self.names
        else:
            c.execute('select %s from %s' % (self.id, self.table_name))
            name_id_pairs = c.fetchall() 
            self.ids = [x[0] for x in name_id_pairs]

            # insert names and ids into table values  
            self.table_vals[self.id] = self.ids
            

        # mark columns as retrieved in status table 
        self.n_rows = len(self.ids)
        status_table = zeros((self.n_rows, len(self.col_names)))

        if self.status_table == None:
            self.status_table = status_table
        else: # re-pulling after adding new row
            if self.status_table.shape[0] == 0:
                self.status_table = status_table
            # change status table shape
            # only ROWS should be added
            if status_table.shape[1] > self.status_table.shape[1]:
                raise Exception("columns have been lost!")

            elif status_table.shape[0] > self.status_table.shape[0]:
                # add rows
                n_new_rows = status_table.shape[0] - self.status_table.shape[0]
                n_cols = self.status_table.shape[1]
                self.status_table = vstack((self.status_table, zeros((n_new_rows, n_cols)))) 

        if self.name_unique:
            self.status_table[:, self.col_names.index(self.id)] = PULLED
            self.status_table[:, self.col_names.index('name')] = PULLED
            
    def __getitem__(self, args):
        if not len(args) == 2:
            raise Exception("lazy_table.__getitem__ takes a 2-tuple as argument")

        row_id = args[0]
        col_id = args[1]
        if isinstance(row_id, list) or isinstance(row_id, tuple):
            return [self.__getitem__((r, col_id)) for r in row_id]

        if not (isinstance(col_id, basestring) or col_id == slice(None)):
            raise Exception("column name must be a string or full slice!")
        elif isinstance(col_id, basestring) and not col_id in self.table_vals.keys():
            return None

        if hasattr(row_id, 'name'):
            row_id = row_id.name

        if isinstance(row_id, int) or isinstance(row_id, long):
            row_id = self.ids.index(row_id)
            return self.lazy_lookup(col_id, row_id)
        elif isinstance(row_id, slice):
            return self.lazy_lookup(col_id, row_id)
        elif isinstance(row_id, basestring) and not self.name_unique:
            raise Exception('Name indexing is not allowed because names are not unique in this table')
        elif isinstance(row_id, basestring):
            # lookup name 
            row_id = self.names.index(row_id)
            return self.lazy_lookup(col_id, row_id)

    def lazy_lookup(self, col_id, row_idx):
        if col_id == slice(None):
            data = {}
            for col in self.col_names:
                data[col] = self.lazy_lookup(col, row_idx)
            return data 

        rows = arange(self.n_rows)[row_idx]
        if not isinstance(rows, ndarray):
            rows = array([rows])
        row_ids = array( [self.ids[row_idx]] ).reshape(-1)

        # check status table
        try:
            col_idx = self.col_names.index(col_id)
        except:
            return None
        status = self.status_table[row_idx, col_idx]

        # determine which rows need to be retrieved 
        if isinstance(status, ndarray):
            retr_row_ids = row_ids[ find(status == UNKNOWN) ]
            retr_rows = rows[ find(status == UNKNOWN) ]
        elif status == UNKNOWN:
            retr_row_ids = array([row_ids]).reshape(-1)
            retr_rows = array([rows]).reshape(-1)
        else:
            retr_rows = []

        if len(retr_rows) > 0:
            # retrieve data from db table
            c = self.conn.cursor()
            val = mysql_get_entry(c, col_id, self.table_name, retr_row_ids, 
                id=self.id)

            for k, row in enumerate(retr_rows):
                self.table_vals[col_id][row] = val[k]
                self.status_table[row, col_idx] = PULLED
        
        ret_val =  self.table_vals[col_id][row_idx]
        if not self.cache_data:
            self.table_vals[col_id][row_idx] = None
        return ret_val

    
    def __setitem__(self, args, val):
        row_id = args[0]
        col_id = args[1]

        if not isinstance(col_id, basestring):
            print "column name must be a string!"
            raise Exception

        if not col_id in self.table_vals.keys(): # new table column 
            # add col name to self.col_names
            self.col_names.append( col_id )

            # add a new column to the table status 
            self.status_table = hstack((self.status_table, 
                MODIFIED*ones((self.n_rows, 1))))            

            # add new elem to table_vals_dict
            self.table_vals[col_id] = [None]*self.n_rows

        # get column index
        col_idx = self.col_names.index(col_id)

        if hasattr(row_id, 'name'):
            row_id = row_id.name

        if isinstance(row_id, basestring) and self.name_unique: # lookup name 
            if not row_id in self.names:
                # insert into table
                c = self.conn.cursor()
                c.execute("insert into %s (name) values ('%s')" % (self.table_name, row_id))

                # update id/name lists  
                self.update_keys()

                # extend the list representing each column of values
                n_rows = len(self.table_vals[self.id])
                for col in self.col_names:
                    self.table_vals[col] += [None]*(n_rows -len(self.table_vals[col]))

            row_id = self.names.index(row_id)
        elif isinstance(row_id, basestring): 
            row_id = len(self.names)
        elif isinstance(row_id, slice):
            row_id = arange(self.n_rows)[row_id]
        elif isinstance(row_id, int) or isinstance(row_id, long):
            if row_id == -1: 
                c = self.conn.cursor()
                if col_id == 'name':
                    c.execute("insert into %s (name) values ('%s')" % (self.table_name, val))
                else:
                    c.execute("insert into %s () values ()" % (self.table_name))

                self.update_keys()

                # extend the list representing each column of values
                n_rows = len(self.table_vals[self.id])
                for col in self.col_names:
                    self.table_vals[col] += [None]*(n_rows -len(self.table_vals[col]))
                row_id = len(self.ids)-1
            else:
                row_id = self.ids.index(row_id)
        elif isinstance(row_id, float):
            raise Exception("Floating point keys are not a good idea..")

        if np.iterable(row_id): 
            for id, v in izip(row_id, val):
                self.table_vals[col_id][id] = v
                self.status_table[id, col_idx] = MODIFIED
        else:
            self.table_vals[col_id][row_id] = val
            self.status_table[row_id, col_idx] = MODIFIED
            
    def load(self, row):
        if 'name' not in row:
            raise Exception("row dictionary must contain 'name'")
        name = row['name']
        if not self.name_unique:
            self[-1, 'name'] = name
            row_idx = self.ids[-1] 
        else:
            row_idx = self.names.index(name)

        for col, value in row.items():
            if col in self.col_names:
                self[row_idx, col] = value

    def add_col(self, col_name, dtype='longblob', mym_data=True):
        if isinstance(col_name, list):
            for col in col_name:
                self.add_col(col, dtype=dtype, mym_data=mym_data)
            return 

        orig_col_names = [x[0] for x in self.table_descr]            
        if col_name in orig_col_names:
            return 
        
        # add new columns to table, as type BLOB, NULL 
        c = self.conn.cursor()
        execute_silent(c, MYSQL_ADD_COL, (self.table_name, col_name, dtype)) 

        # Update lazy representation of column names
        self.update_column_names()
        
        #col_id = self.col_names.index(col_name)
        self.table_vals[col_name] = [None]*self.n_rows

        # TODO..the next code assumes that the column added is the last column (which is the default) 
        n_rows = self.status_table.shape[0]
        self.status_table = np.hstack([self.status_table, np.ones((n_rows,1))])

    def delete_rows_with_name(self, name):
        assert 'name' in self.col_names
        names = self[:,'name']
        rows = filter(lambda x: x[1] == name, zip(self.ids, names)) 
        for id, name in rows:
            self.drop_row(id)
            

    def drop_col(self, col):
        if isinstance(col):
            [self.drop_col(c) for c in col]
            return
        elif not col in self.col_names:
            raise Exception("Unknown column being deleted: %s" % col)
        col_index = self.col_names.index(col)

        c = self.conn.cursor()
        c.execute(MYSQL_DROP_COL % (self.table_name, col) )

        # delete column in status table
        np.delete(self.status_table, col_index, axis=1)

        # delete data
        del self.table_vals[col]

    def drop_row(self, row):
        if isinstance(row, basestring) and not self.name_unique:
            raise Exception('Name indexing is not allowed because names are not unique in this table')
        elif isinstance(row, basestring):
            row = self.ids[self.names.index(row)]

        try:
            row = int(row)
        except:
            raise Exception("Row cannot be translated to integer id!")

        c = self.conn.cursor()
        c.execute('delete from %s where %s=%d' % (self.table_name, self.id, row))

        ## Delete lazy representations
        row_idx = self.ids.index(row)
        for col in self.col_names:
            del self.table_vals[col][row_idx]

        self.status_table = np.delete(self.status_table, row_idx, axis=0)
        
    def save(self, verbose=False):
        c = self.conn.cursor()

        # add new columns to MySQL
        orig_col_names = [x[0] for x in self.table_descr]            
        for col in self.col_names:
            if col not in orig_col_names:
                c.execute(MYSQL_ADD_COL %(self.table_name, col, 'longblob')) 
        self.update_column_names()

        # insert individual values into table
        for k, col in enumerate(self.col_names):
            # find inds for values to be inserted for this column 
            mod_data_inds = find(self.status_table[:, k] == MODIFIED)
            ids = array(self.ids)[mod_data_inds] 

            for ind, id in izip(mod_data_inds, ids):
                val = self.table_vals[col][ind]
                use_mym = self.table_descr[k][1] == 'longblob'
                if use_mym: 
                    val = MySQLdb.escape_string(pymym.pack(val))
                else:
                    val = MySQLdb.escape_string(val)

                sql_cmd = MYSQL_UPDATE_ROW % (self.table_name, col, val, self.id, int(id))
                c.execute(sql_cmd)
                if verbose:
                    print sql_cmd
                self.status_table[ ind, k ] = PULLED

        c.execute('describe %s' % self.table_name)
        self.table_descr = c.fetchall()
        self.col_names =[x[0] for x in self.table_descr] 
        
    def flush(self, col_name):
        """ Erase column from table """
        c = self.conn.cursor()
        c.execute("update %s set %s = NULL" % (self.table_name, col_name))

    def search(self, *args, **kwargs):
        """ Find a record specified in kwargs. Uses introspection
        """
        db_string_values = []
        for col, val in kwargs.items():
            if col not in self.col_names:
                continue
            if col == 'name':
                pass
            else:
                val = pymym.pack(val)
            db_string_values.append((col, MySQLdb.escape_string(val)))

        search_params = ' and '.join(["%s='%s'" % (c,v) for c,v in db_string_values])
        #(col, MySQLdb.escape_string(pymym.pack(val))) for col, val in kwargs.items()])
        c = self.conn.cursor()
        sql_cmd = "select id from %s where %s" % (self.table_name, search_params)
        try:
            c.execute(sql_cmd)
        except:
            print sql_cmd
            raise Exception("SQL command execution error")
        try:
            return list(c.fetchall()[0])
        except:
            return []


    def __contains__(self, other):  
        if isinstance(other, basestring):
            return other in self.names
        else:
            raise Exception("__contains__ only defined for string arguments")

    def __str__(self):
        return "Table: %s\nDatabase: %s\nHost: %s" % (self.table_name, self.db_name, self.conn.get_host_info()) 

    def __repr__(self):
        return self.__str__()

    def get_index(self, name):
        """ Return the row index of the record with the given name """
        try:
            return self.ids[self.names.index(name)]
        except:
            return None

    def operator(self, fn):
        for k, name in izip(self.ids, self.names):
            try:
                fn(name)
            except:
                pass
        self.save()


            
def check_tbl_for_value(val):
    """ Return function to check table for value 
    """
    return lambda s, table: table[s, val] == '' or table[s,val] == None



